import gym
import torch
from PIL import Image

from 主要 import 记忆体, 近端策略优化


def 主要():
    游戏环境名 = "LunarLander-v2"
    游戏环境 = gym.make(游戏环境名, render_mode='human')
    状态_维度 = 游戏环境.observation_space.shape[0]
    动作_维度 = 4
    是否渲染 = True
    """
        时间步长：
        许多物理方程将以这样的形式出现：“从现在开始 1 秒的速度 = 现在的速度 + 1 秒的加速度将产生的速度”（并且大多数不适合这种形式的方程可以以某种方式转换...... )
        “1 秒”部分是时间步长。定期改变世界的游戏会执行大量此类计算，并且要保持正常运行，需要使用一致的时间步长值评估所有这些物理方程。
        游戏经常使用接近每个渲染图形帧之间周期的时间步长，或图形帧之间的理想时间，例如 0.01667 秒，但这不是必需的。一些像 BeamNG.drive 这样的游戏会比游戏渲染更频繁地​​计算物理，以生成更平滑的物理近似。
        其他游戏将不那么频繁地计算物理并混合每个点之间渲染的位置。他们的方法将取决于每个游戏。
    """
    时间步长_最大值 = 300  # 时间的步长，这是一个单位
    隐藏层_变量数 = 64  # 隐藏层的变量数
    时间步长_更新_间隔 = 2000
    学习率 = 0.0007
    贝塔值 = (0.9, 0.999)  # 用来计算梯度的平均数和平方的系数
    伽马值 = 0.99  # 折扣因子
    轮回次数 = 4  # 模型更新策略
    周期数 = 3
    是否保存动图 = False

    # 文件夹中保存的模型名称可能出现乱码，需要修改名称
    文件名 = '近端策略优化_{}.pth'.format(游戏环境名)

    夹子_间距 = 0.2  # 夹子函数的范围

    一个记忆体 = 记忆体()

    一个近端策略优化 = 近端策略优化(状态_维度, 动作_维度, 隐藏层_变量数, 学习率, 贝塔值, 伽马值, 轮回次数, 夹子_间距)
    一个近端策略优化.旧策略.load_state_dict(torch.load(文件名), strict=False)

    for 索引_周期 in range(1, 周期数 + 1):
        周期_奖励 = 0
        状态, _ = 游戏环境.reset()
        for t in range(时间步长_最大值):
            动作 = 一个近端策略优化.旧策略.动作(状态,一个记忆体)
            状态, 奖励, 完毕, _, _ = 游戏环境.step(动作)
            周期_奖励 += 奖励
            if 是否渲染:
                游戏环境.render()
            if 是否保存动图:
                游戏环境.render_mode = 'rgb_array'
                图片 = 游戏环境.render()
                图片 = Image.fromarray(图片)
                图片.save("动图_{}.gif".format(t))
            if 完毕:
                break
        print("周期：{} \t奖励：{}".format(索引_周期, 周期_奖励))
        周期_奖励 = 0
    游戏环境.close()


if __name__ == '__main__':
    主要()
